# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
import asyncio
from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union

from camel.interpreters.base import BaseInterpreter
from camel.interpreters.interpreter_error import InterpreterError
from camel.logger import get_logger

logger = get_logger(__name__)


class MicrosandboxInterpreter(BaseInterpreter):
    r"""Microsandbox Code Interpreter implementation.

    This interpreter provides secure code execution using microsandbox,
    a self-hosted platform for secure execution of untrusted user/AI code.
    It supports Python code execution via PythonSandbox, JavaScript/Node.js
    code execution via NodeSandbox, and shell commands via the command
    interface.

    Args:
        require_confirm (bool, optional): If True, prompt user before running
            code strings for security. (default: :obj:`True`)
        server_url (str, optional): URL of the microsandbox server. If not
            provided, will use MSB_SERVER_URL environment variable, then
            fall back to http://127.0.0.1:5555. (default: :obj:`None`)
        api_key (str, optional): API key for microsandbox authentication.
            If not provided, will use MSB_API_KEY environment variable.
            (default: :obj:`None`)
        namespace (str, optional): Namespace for the sandbox.
            (default: :obj:`"default"`)
        sandbox_name (str, optional): Name of the sandbox instance. If not
            provided, a random name will be generated by the SDK.
            (default: :obj:`None`)
        timeout (int, optional): Default timeout for code execution in seconds.
            (default: :obj:`30`)

    Environment Variables:
        MSB_SERVER_URL: URL of the microsandbox server.
        MSB_API_KEY: API key for microsandbox authentication.

    Note:
        The SDK handles parameter priority as: user parameter > environment
            variable > default value.
    """

    _CODE_TYPE_MAPPING: ClassVar[Dict[str, str]] = {
        # Python code - uses PythonSandbox
        "python": "python_sandbox",
        "py3": "python_sandbox",
        "python3": "python_sandbox",
        "py": "python_sandbox",
        # JavaScript/Node.js code - uses NodeSandbox
        "javascript": "node_sandbox",
        "js": "node_sandbox",
        "node": "node_sandbox",
        "typescript": "node_sandbox",
        "ts": "node_sandbox",
        # Shell commands - uses command.run()
        "bash": "shell_command",
        "shell": "shell_command",
        "sh": "shell_command",
    }

    def __init__(
        self,
        require_confirm: bool = True,
        server_url: Optional[str] = None,
        api_key: Optional[str] = None,
        namespace: str = "default",
        sandbox_name: Optional[str] = None,
        timeout: int = 30,
    ) -> None:
        from microsandbox import (
            NodeSandbox,
            PythonSandbox,
        )

        # Store parameters, let SDK handle defaults and environment variables
        self.require_confirm = require_confirm
        self.server_url = server_url  # None means use SDK default logic
        self.api_key = api_key  # None means use SDK default logic
        self.namespace = namespace
        self.sandbox_name = (
            sandbox_name  # None means SDK generates random name
        )
        self.timeout = timeout

        # Store sandbox configuration
        self._sandbox_config = {
            "server_url": self.server_url,
            "namespace": self.namespace,
            "name": self.sandbox_name,
            "api_key": self.api_key,
        }

        # Store sandbox classes for reuse
        self._PythonSandbox = PythonSandbox
        self._NodeSandbox = NodeSandbox

        # Log initialization info
        logger.info("Initialized MicrosandboxInterpreter")
        logger.info(f"Namespace: {self.namespace}")
        if self.sandbox_name:
            logger.info(f"Sandbox name: {self.sandbox_name}")
        else:
            logger.info("Sandbox name: will be auto-generated by SDK")

    def run(
        self,
        code: str,
        code_type: str = "python",
    ) -> str:
        r"""Executes the given code in the microsandbox.

        Args:
            code (str): The code string to execute.
            code_type (str): The type of code to execute. Supported types:
                'python', 'javascript', 'bash'. (default: :obj:`python`)

        Returns:
            str: The string representation of the output of the executed code.

        Raises:
            InterpreterError: If the `code_type` is not supported or if any
                runtime error occurs during the execution of the code.
        """
        if code_type not in self._CODE_TYPE_MAPPING:
            raise InterpreterError(
                f"Unsupported code type {code_type}. "
                f"`{self.__class__.__name__}` only supports "
                f"{', '.join(list(self._CODE_TYPE_MAPPING.keys()))}."
            )

        # Print code for security checking
        if self.require_confirm:
            logger.info(
                f"The following {code_type} code will run on "
                f"microsandbox: {code}"
            )
            self._confirm_execution("code")

        # Run the code asynchronously
        return asyncio.run(self._run_async(code, code_type))

    async def _run_async(self, code: str, code_type: str) -> str:
        r"""Asynchronously executes code in microsandbox.

        Args:
            code (str): The code to execute.
            code_type (str): The type of code to execute.

        Returns:
            str: The output of the executed code.

        Raises:
            InterpreterError: If execution fails.
        """
        try:
            execution_method = self._CODE_TYPE_MAPPING[code_type]

            if execution_method == "python_sandbox":
                return await self._run_python_code(code)
            elif execution_method == "node_sandbox":
                return await self._run_node_code(code)
            elif execution_method == "shell_command":
                return await self._run_shell_command(code)
            else:
                raise InterpreterError(
                    f"Unsupported execution method: {execution_method}"
                )

        except Exception as e:
            raise InterpreterError(
                f"Error executing code in microsandbox: {e}"
            )

    async def _run_python_code(self, code: str) -> str:
        r"""Execute Python code using PythonSandbox.

        Args:
            code (str): Python code to execute.

        Returns:
            str: Execution output.
        """
        async with self._PythonSandbox.create(
            **self._sandbox_config
        ) as sandbox:
            execution = await asyncio.wait_for(
                sandbox.run(code), timeout=self.timeout
            )
            return await self._get_execution_output(execution)

    async def _run_node_code(self, code: str) -> str:
        r"""Execute JavaScript/Node.js code using NodeSandbox.

        Args:
            code (str): JavaScript/Node.js code to execute.

        Returns:
            str: Execution output.
        """
        async with self._NodeSandbox.create(**self._sandbox_config) as sandbox:
            execution = await asyncio.wait_for(
                sandbox.run(code), timeout=self.timeout
            )
            return await self._get_execution_output(execution)

    async def _run_shell_command(self, code: str) -> str:
        r"""Execute shell commands directly.

        Args:
            code (str): Shell command to execute.

        Returns:
            str: Command execution output.
        """
        # Use any sandbox for shell commands
        async with self._PythonSandbox.create(
            **self._sandbox_config
        ) as sandbox:
            execution = await asyncio.wait_for(
                sandbox.command.run("bash", ["-c", code]), timeout=self.timeout
            )
            return await self._get_command_output(execution)

    async def _get_execution_output(self, execution) -> str:
        r"""Get output from code execution.

        Args:
            execution: Execution object from sandbox.run().

        Returns:
            str: Formatted execution output.
        """
        output = await execution.output()
        error = await execution.error()

        result_parts = []
        if output and output.strip():
            result_parts.append(output.strip())
        if error and error.strip():
            result_parts.append(f"STDERR: {error.strip()}")

        return (
            "\n".join(result_parts)
            if result_parts
            else "Code executed successfully (no output)"
        )

    async def _get_command_output(self, execution) -> str:
        r"""Get output from command execution.

        Args:
            execution: CommandExecution object from sandbox.command.run().

        Returns:
            str: Formatted command output.
        """
        output = await execution.output()
        error = await execution.error()

        result_parts = []
        if output and output.strip():
            result_parts.append(output.strip())
        if error and error.strip():
            result_parts.append(f"STDERR: {error.strip()}")
        if hasattr(execution, 'exit_code') and execution.exit_code != 0:
            result_parts.append(f"Exit code: {execution.exit_code}")

        return (
            "\n".join(result_parts)
            if result_parts
            else "Command executed successfully (no output)"
        )

    def _confirm_execution(self, execution_type: str) -> None:
        r"""Prompt user for confirmation before executing code or commands.

        Args:
            execution_type (str): Type of execution ('code' or 'command').

        Raises:
            InterpreterError: If user declines to run the code/command.
        """
        while True:
            choice = input(f"Running {execution_type}? [Y/n]:").lower()
            if choice in ["y", "yes", "ye"]:
                break
            elif choice not in ["no", "n"]:
                continue
            raise InterpreterError(
                f"Execution halted: User opted not to run the "
                f"{execution_type}. "
                f"This choice stops the current operation and any "
                f"further {execution_type} execution."
            )

    def supported_code_types(self) -> List[str]:
        r"""Provides supported code types by the interpreter."""
        return list(self._CODE_TYPE_MAPPING.keys())

    def update_action_space(self, action_space: Dict[str, Any]) -> None:
        r"""Updates action space for interpreter.

        Args:
            action_space: Action space dictionary (unused in microsandbox).

        Note:
            Microsandbox doesn't support action space updates as it runs
            in isolated environments for each execution.
        """
        # Explicitly acknowledge the parameter to avoid linting warnings
        _ = action_space
        logger.warning(
            "Microsandbox doesn't support action space updates. "
            "Code runs in isolated environments for each execution."
        )

    def execute_command(self, command: str) -> Union[str, Tuple[str, str]]:
        r"""Execute a shell command in the microsandbox.

        This method is designed for package management and system
        administration tasks. It executes shell commands directly
        using the microsandbox command interface.

        Args:
            command (str): The shell command to execute (e.g.,
            "pip install numpy", "ls -la", "apt-get update").

        Returns:
            Union[str, Tuple[str, str]]: The output of the command.

        Examples:
            >>> interpreter.execute_command("pip install numpy")
            >>> interpreter.execute_command("npm install express")
            >>> interpreter.execute_command("ls -la /tmp")
        """
        # Print command for security checking
        if self.require_confirm:
            logger.info(
                f"The following shell command will run on "
                f"microsandbox: {command}"
            )
            self._confirm_execution("command")

        return asyncio.run(self._execute_command_async(command))

    async def _execute_command_async(self, command: str) -> str:
        r"""Asynchronously executes a shell command in microsandbox.

        Args:
            command (str): The shell command to execute.

        Returns:
            str: The output of the command execution.

        Raises:
            InterpreterError: If execution fails.
        """
        try:
            async with self._PythonSandbox.create(
                **self._sandbox_config
            ) as sandbox:
                execution = await asyncio.wait_for(
                    sandbox.command.run("bash", ["-c", command]),
                    timeout=self.timeout,
                )
                return await self._get_command_output(execution)

        except Exception as e:
            raise InterpreterError(
                f"Error executing command in microsandbox: {e}"
            )

    def __del__(self) -> None:
        r"""Destructor for the MicrosandboxInterpreter class.

        Microsandbox uses context managers for resource management,
        so no explicit cleanup is needed.
        """
        logger.debug("MicrosandboxInterpreter cleaned up")
